import os
import random
import time
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import entropy

class SecurityTester:
    """
    安全性测试类
    包含对加密算法进行安全性评估的各种方法
    """
    
    def __init__(self):
        pass
    
    def run_nist_randomness_tests(self, data):
        """
        执行简化版的NIST随机性测试
        :param data: 被测试的字节数据
        :return: 测试结果字典
        """
        # 将字节数据转换为位序列
        bit_seq = ''.join(format(byte, '08b') for byte in data)
        bits = np.array([int(bit) for bit in bit_seq])
        
        results = {}
        
        # 1. 频率（单比特）测试
        results['frequency_test'] = self._frequency_test(bits)
        
        # 2. 游程测试
        results['runs_test'] = self._runs_test(bits)
        
        # 3. 块频率测试
        results['block_frequency_test'] = self._block_frequency_test(bits, block_size=128)
        
        # 4. 序列测试
        results['serial_test'] = self._serial_test(bits, block_size=2)
        
        # 5. 累积和测试
        results['cumulative_sums_test'] = self._cumulative_sums_test(bits)
        
        # 汇总结果
        passed_tests = sum(1 for result in results.values() if result['pass'])
        total_tests = len(results)
        
        results['summary'] = {
            'passed_tests': passed_tests,
            'total_tests': total_tests,
            'pass_rate': passed_tests / total_tests if total_tests > 0 else 0,
            'overall_pass': passed_tests == total_tests
        }
        
        return results
    
    def _frequency_test(self, bits):
        """
        频率（单比特）测试
        测试序列中0和1的比例是否接近1:1
        """
        n = len(bits)
        ones_count = np.sum(bits)
        zeros_count = n - ones_count
        
        # 计算0和1的比例
        proportion = ones_count / n
        
        # 理想情况下，比例应接近0.5
        threshold = 0.1  # 允许的偏差
        passed = abs(proportion - 0.5) < threshold
        
        return {
            'name': '频率测试',
            'proportion_ones': proportion,
            'proportion_zeros': zeros_count / n,
            'pass': passed
        }
    
    def _runs_test(self, bits):
        """
        游程测试
        测试0和1的连续出现的游程数是否符合随机序列的特性
        """
        n = len(bits)
        runs = 1  # 初始游程数
        
        # 计算游程数
        for i in range(1, n):
            if bits[i] != bits[i-1]:
                runs += 1
        
        # 计算期望游程数（对随机序列，约为n/2）
        expected_runs = n / 2
        
        # 允许的偏差
        threshold = 0.2 * expected_runs
        passed = abs(runs - expected_runs) < threshold
        
        return {
            'name': '游程测试',
            'runs': runs,
            'expected_runs': expected_runs,
            'pass': passed
        }
    
    def _block_frequency_test(self, bits, block_size=128):
        """
        块频率测试
        将位序列分成多个块，测试每个块中1的比例
        """
        n = len(bits)
        num_blocks = n // block_size
        
        if num_blocks == 0:
            return {
                'name': '块频率测试',
                'error': '数据长度不足',
                'pass': False
            }
        
        proportions = []
        for i in range(num_blocks):
            block = bits[i * block_size:(i + 1) * block_size]
            proportion = np.sum(block) / block_size
            proportions.append(proportion)
        
        # 计算每个块中1的比例的标准差
        std_dev = np.std(proportions)
        
        # 理想情况下，标准差应该较小
        threshold = 0.15
        passed = std_dev < threshold
        
        return {
            'name': '块频率测试',
            'std_dev': std_dev,
            'avg_proportion': np.mean(proportions),
            'pass': passed
        }
    
    def _serial_test(self, bits, block_size=2):
        """
        序列测试
        测试所有可能的长度为m的模式出现的频率是否接近均匀分布
        """
        n = len(bits)
        patterns = {}
        
        # 计算所有可能的模式
        for i in range(n - block_size + 1):
            pattern = ''.join(str(bit) for bit in bits[i:i+block_size])
            patterns[pattern] = patterns.get(pattern, 0) + 1
        
        # 理论上每种模式的期望出现次数
        expected_count = (n - block_size + 1) / (2 ** block_size)
        
        # 计算卡方统计量
        chi_square = 0
        for count in patterns.values():
            chi_square += ((count - expected_count) ** 2) / expected_count
        
        # 对于块大小为2，自由度为3，95%置信度的卡方临界值约为7.815
        threshold = 7.815
        passed = chi_square < threshold
        
        return {
            'name': '序列测试',
            'chi_square': chi_square,
            'expected_count': expected_count,
            'pattern_counts': patterns,
            'pass': passed
        }
    
    def _cumulative_sums_test(self, bits):
        """
        累积和测试
        测试序列累积和的最大偏差
        """
        n = len(bits)
        # 将0转换为-1，计算累积和
        mapped_bits = np.where(bits == 0, -1, 1)
        cumulative_sum = np.cumsum(mapped_bits)
        
        # 计算最大偏差
        max_deviation = max(abs(cumulative_sum))
        
        # 理想情况下，最大偏差应该较小
        # 根据随机游走理论，最大偏差的期望约为sqrt(n)
        expected_max_deviation = np.sqrt(n)
        threshold = 3 * expected_max_deviation  # 允许的最大偏差
        
        passed = max_deviation < threshold
        
        return {
            'name': '累积和测试',
            'max_deviation': max_deviation,
            'expected_max_deviation': expected_max_deviation,
            'pass': passed
        }
    
    def simulate_known_plaintext_attack(self, encrypt_func, key, num_samples=100):
        """
        模拟已知明文攻击
        :param encrypt_func: 加密函数，接受(plaintext, key)参数
        :param key: 加密密钥
        :param num_samples: 样本数量
        :return: 攻击结果
        """
        plaintext_samples = []
        ciphertext_samples = []
        
        # 生成随机明文样本
        for _ in range(num_samples):
            # 生成随机长度的明文
            length = random.randint(16, 64)
            plaintext = os.urandom(length)
            plaintext_samples.append(plaintext)
            
            # 加密
            ciphertext = encrypt_func(plaintext, key)
            ciphertext_samples.append(ciphertext)
        
        # 分析密文与明文的统计关系
        correlations = []
        
        # 计算每个样本对的明文和密文字节之间的相关性
        for p, c in zip(plaintext_samples, ciphertext_samples):
            min_len = min(len(p), len(c))
            p_bytes = np.array([b for b in p[:min_len]])
            c_bytes = np.array([b for b in c[:min_len]])
            
            if min_len > 1:
                # 计算相关系数
                correlation = np.corrcoef(p_bytes, c_bytes)[0, 1]
                if not np.isnan(correlation):
                    correlations.append(correlation)
        
        # 平均相关系数
        avg_correlation = np.mean(correlations) if correlations else 0
        
        # 计算信息熵
        plaintext_entropy = np.mean([self._calculate_entropy(p) for p in plaintext_samples])
        ciphertext_entropy = np.mean([self._calculate_entropy(c) for c in ciphertext_samples])
        
        # 评估结果
        is_secure = avg_correlation < 0.1 and ciphertext_entropy > plaintext_entropy * 0.9
        
        return {
            'name': '已知明文攻击模拟',
            'samples': num_samples,
            'avg_correlation': avg_correlation,
            'plaintext_entropy': plaintext_entropy,
            'ciphertext_entropy': ciphertext_entropy,
            'pass': is_secure
        }
    
    def simulate_chosen_ciphertext_attack(self, decrypt_func, key, num_samples=100):
        """
        模拟选择密文攻击
        :param decrypt_func: 解密函数，接受(ciphertext, key)参数
        :param key: 解密密钥
        :param num_samples: 样本数量
        :return: 攻击结果
        """
        ciphertext_samples = []
        decrypted_samples = []
        
        # 生成随机密文样本
        for _ in range(num_samples):
            # 生成随机长度的"密文"（实际上是随机字节序列）
            length = random.randint(32, 128)  # 密文通常比明文长
            ciphertext = os.urandom(length)
            ciphertext_samples.append(ciphertext)
            
            try:
                # 尝试解密（可能会失败，因为随机字节不一定是有效的密文）
                decrypted = decrypt_func(ciphertext, key)
                decrypted_samples.append(decrypted)
            except:
                # 解密失败，添加空字节作为占位符
                decrypted_samples.append(b'')
        
        # 过滤掉解密失败的样本
        valid_samples = [(c, d) for c, d in zip(ciphertext_samples, decrypted_samples) if d]
        
        if not valid_samples:
            return {
                'name': '选择密文攻击模拟',
                'samples': num_samples,
                'valid_samples': 0,
                'pass': True,  # 没有有效样本意味着算法对随机输入的健壮性好
                'message': '所有随机密文均无法被成功解密，表明算法对无效输入有良好的抵抗性'
            }
        
        # 分析有效样本
        correlations = []
        for c, d in valid_samples:
            min_len = min(len(c), len(d))
            if min_len > 1:
                c_bytes = np.array([b for b in c[:min_len]])
                d_bytes = np.array([b for b in d[:min_len]])
                
                # 计算相关系数
                correlation = np.corrcoef(c_bytes, d_bytes)[0, 1]
                if not np.isnan(correlation):
                    correlations.append(correlation)
        
        # 平均相关系数
        avg_correlation = np.mean(correlations) if correlations else 0
        
        # 有效样本比例
        valid_ratio = len(valid_samples) / num_samples
        
        # 评估结果
        # 通常，有效样本比例低且相关性低表示安全性好
        is_secure = valid_ratio < 0.1 and avg_correlation < 0.1
        
        return {
            'name': '选择密文攻击模拟',
            'samples': num_samples,
            'valid_samples': len(valid_samples),
            'valid_ratio': valid_ratio,
            'avg_correlation': avg_correlation,
            'pass': is_secure
        }
    
    def analyze_forward_security(self, encrypt_func, rotate_key_func, plaintext, num_rotations=5):
        """
        评估前向安全性
        :param encrypt_func: 加密函数，接受(plaintext, key)参数
        :param rotate_key_func: 密钥轮换函数，接受(key)参数并返回新密钥
        :param plaintext: 用于测试的明文
        :param num_rotations: 密钥轮换次数
        :return: 分析结果
        """
        ciphertexts = []
        keys = []
        
        # 生成初始密钥
        current_key = os.urandom(32)  # 假设使用256位密钥
        keys.append(current_key)
        
        # 使用初始密钥加密
        ciphertext = encrypt_func(plaintext, current_key)
        ciphertexts.append(ciphertext)
        
        # 执行多次密钥轮换并加密相同的明文
        for _ in range(num_rotations):
            current_key = rotate_key_func(current_key)
            keys.append(current_key)
            
            ciphertext = encrypt_func(plaintext, current_key)
            ciphertexts.append(ciphertext)
        
        # 分析不同密钥产生的密文之间的相似性
        similarities = []
        for i in range(len(ciphertexts)):
            for j in range(i + 1, len(ciphertexts)):
                ci = ciphertexts[i]
                cj = ciphertexts[j]
                
                # 计算汉明距离的近似值（XOR后的1的数量）
                min_len = min(len(ci), len(cj))
                ci_bytes = np.array([b for b in ci[:min_len]])
                cj_bytes = np.array([b for b in cj[:min_len]])
                
                # 计算归一化的汉明距离
                xor_result = np.bitwise_xor(ci_bytes, cj_bytes)
                ones_count = np.sum([bin(b).count('1') for b in xor_result])
                total_bits = min_len * 8
                normalized_distance = ones_count / total_bits
                
                # 相似性 = 1 - 归一化距离
                similarity = 1 - normalized_distance
                similarities.append(similarity)
        
        # 平均相似性
        avg_similarity = np.mean(similarities) if similarities else 0
        
        # 评估结果
        # 理想情况下，不同密钥加密的相同明文应该产生完全不同的密文
        # 因此平均相似性应接近0.5（随机比特序列的期望值）
        is_secure = abs(avg_similarity - 0.5) < 0.1
        
        return {
            'name': '前向安全性分析',
            'rotations': num_rotations,
            'avg_similarity': avg_similarity,
            'expected_similarity': 0.5,
            'pass': is_secure
        }
    
    def _calculate_entropy(self, data):
        """计算字节序列的香农熵"""
        # 统计每个字节值的频率
        freq = np.zeros(256, dtype=float)
        for byte in data:
            freq[byte] += 1
        
        # 计算熵
        freq = freq[freq > 0] / len(data)  # 概率分布
        return entropy(freq, base=2)
    
    def visualize_distribution(self, data, title="字节分布"):
        """
        可视化字节分布
        :param data: 字节数据
        :param title: 图表标题
        """
        # 统计每个字节值的频率
        freq = np.zeros(256, dtype=int)
        for byte in data:
            freq[byte] += 1
        
        # 绘制直方图
        plt.figure(figsize=(12, 6))
        plt.bar(range(256), freq, width=1.0)
        plt.title(title)
        plt.xlabel('字节值')
        plt.ylabel('频率')
        plt.xlim(0, 255)
        plt.grid(axis='y', alpha=0.75)
        
        return plt
    
    def visualize_entropy(self, original_data, encrypted_data, block_size=1024):
        """
        可视化原始数据和加密数据的局部熵
        :param original_data: 原始字节数据
        :param encrypted_data: 加密字节数据
        :param block_size: 块大小
        """
        # 计算分块熵
        def calculate_block_entropy(data, size):
            blocks = [data[i:i+size] for i in range(0, len(data), size)]
            return [self._calculate_entropy(block) for block in blocks if block]
        
        original_entropy = calculate_block_entropy(original_data, block_size)
        encrypted_entropy = calculate_block_entropy(encrypted_data, block_size)
        
        # 确保两个列表长度相同（取最小值）
        min_len = min(len(original_entropy), len(encrypted_entropy))
        original_entropy = original_entropy[:min_len]
        encrypted_entropy = encrypted_entropy[:min_len]
        
        # 绘制比较图
        plt.figure(figsize=(12, 6))
        plt.plot(range(min_len), original_entropy, 'b-', label='原始数据')
        plt.plot(range(min_len), encrypted_entropy, 'r-', label='加密数据')
        plt.title('数据块熵比较')
        plt.xlabel('数据块索引')
        plt.ylabel('熵 (bits)')
        plt.legend()
        plt.grid(True, alpha=0.3)
        
        return plt 